import torch
from torch import Tensor

from config import device
from problems.LSMOP.LSMOP_1 import LSMOP1
from problems.problem import Problem


class SBX:
    @staticmethod
    def do(parent0: Tensor, parent1: Tensor, problem: Problem, offspring_size: int = 1, pro_c: float = 0.9,
           dis_c: float = 20) -> Tensor:
        """
        模拟二进制交叉
        :param parent0:父代的dec
        :param parent1:父代的dec
        :param problem:优化问题
        :param offspring_size: 产生个后代个数
        :param pro_c: 交叉概率
        :param dis_c: n越大，产生的个体和父代越相似
        :return: 返回后代个体
        """
        if parent0.dim() == 1:
            parent0 = parent0.unsqueeze(dim=0)

        if parent1.dim() == 1:
            parent1 = parent1.unsqueeze(dim=0)

        var_dim = problem.var_dim
        parent0_size = parent0.size(0)
        parent1_size = parent1.size(0)
        beta = torch.zeros(parent0_size, var_dim, device=device)
        mu = torch.rand(parent0_size, var_dim, device=device)

        beta[mu <= 0.5] = (2 * mu[mu <= 0.5]) ** (1 / (dis_c + 1))
        beta[mu > 0.5] = (2 - 2 * mu[mu > 0.5]) ** (-1 / (dis_c + 1))
        beta = beta * (-1) ** torch.randint(2, (parent0_size, var_dim), device=device)
        beta[torch.rand(parent0_size, var_dim, device=device) < 0.5] = 1
        beta[torch.rand(parent0_size, device=device) > pro_c] = 1
        offspring_list = []

        for i in range(0, parent0_size):
            parent1_temp = parent1 if parent1_size == 1 else parent1[i].unsqueeze(dim=0)
            if offspring_size == 1:
                if torch.rand(1, device=device) < 0.5:
                    offspring = (parent0[i] + parent1_temp) / 2 + beta[i] * (parent0[i] - parent1_temp) / 2
                else:
                    offspring = (parent0[i] + parent1_temp) / 2 - beta[i] * (parent0[i] - parent1_temp) / 2
            elif offspring_size == 2:
                offspring = torch.zeros(2, var_dim, dtype=float, device=device)
                offspring[0] = (parent0[i] + parent1_temp) / 2 + beta[i] * (parent0[i] - parent1_temp) / 2
                offspring[1] = (parent0[i] + parent1_temp) / 2 - beta[i] * (parent0[i] - parent1_temp) / 2
            offspring_list.append(offspring)
        return problem.repair_decision(torch.cat(offspring_list, dim=0))


if __name__ == '__main__':
    print(SBX.do(torch.tensor([1, 2, 3], dtype=float), torch.tensor([2, 1, 0], dtype=float), problem=LSMOP1(3, 2),
                 offspring_size=2))
